#!/usr/bin/env python3
"""
Milvus 快速入门 Demo
简化版本，适合初学者快速上手
"""

import random
from pymilvus import connections, utility, FieldSchema, CollectionSchema, DataType, Collection

def main():
    print("🚀 Milvus 快速入门")
    print("=" * 30)
    
    # 1. 连接Milvus
    print("1️⃣ 连接到Milvus...")
    try:
        connections.connect("default", host="localhost", port="19530")
        print("   ✅ 连接成功!")
    except Exception as e:
        print(f"   ❌ 连接失败: {e}")
        print("   请确保Milvus已启动: docker-compose up -d")
        return
    
    # 2. 创建集合
    print("\n2️⃣ 创建集合...")
    collection_name = "quick_demo"
    
    # 删除已存在的集合
    if utility.has_collection(collection_name):
        utility.drop_collection(collection_name)
    
    # 定义字段
    fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name="name", dtype=DataType.VARCHAR, max_length=100),
        FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=8)  # 使用小维度便于理解
    ]
    
    schema = CollectionSchema(fields, "快速演示集合")
    collection = Collection(collection_name, schema)
    
    # 创建索引
    collection.create_index("vector", {"metric_type": "L2", "index_type": "FLAT"})
    print("   ✅ 集合创建成功!")
    
    # 3. 插入数据
    print("\n3️⃣ 插入示例数据...")
    names = ["Alice", "Bob", "Charlie", "Diana", "Eve"]
    vectors = [[random.random() for _ in range(8)] for _ in range(5)]
    
    collection.insert([names, vectors])
    collection.flush()
    print("   ✅ 插入5条数据!")
    
    # 4. 搜索向量
    print("\n4️⃣ 执行向量搜索...")
    collection.load()
    
    # 搜索最相似的向量
    search_vector = [random.random() for _ in range(8)]
    results = collection.search(
        [search_vector],
        "vector",
        {"metric_type": "L2"},
        limit=3,
        output_fields=["name"]
    )
    
    print("   🔍 搜索结果:")
    for i, hit in enumerate(results[0]):
        print(f"   第{i+1}名: {hit.entity.get('name')} (距离: {hit.distance:.4f})")
    
    # 5. 查询数据
    print("\n5️⃣ 查询所有数据...")
    all_data = collection.query(expr="", output_fields=["id", "name"])
    print("   📊 所有数据:")
    for item in all_data:
        print(f"   ID: {item['id']}, 姓名: {item['name']}")
    
    # 6. 清理
    print("\n6️⃣ 清理资源...")
    utility.drop_collection(collection_name)
    connections.disconnect("default")
    print("   ✅ 清理完成!")
    
    print("\n🎉 快速入门完成!")
    print("💡 提示: 查看 hello_milvus.py 获取更详细的功能演示")

if __name__ == "__main__":
    main() 